import os
import json
from typing import Optional, Dict, Any, List
from alibabacloud_tea_openapi import models as open_api_models
from alibabacloud_bailian20231229.client import Client as BailianClient
from alibabacloud_credentials.models import Config as CredentialConfig
from alibabacloud_bailian20231229 import models as bailian_models
from alibabacloud_credentials.client import Client as CredentialClient
from configs.other_config import ACCESS_KEY_ID, ACCESS_KEY_SECRET, WORKSPACE_ID
import time

class BailianConfig:
    """阿里云百炼服务配置类"""
    def __init__(
        self,
        access_key_id: Optional[str] = None,
        access_key_secret: Optional[str] = None,
        workspace_id: Optional[str] = None,
        endpoint: str = "bailian.cn-beijing.aliyuncs.com"
    ):
        """
        初始化配置
        :param access_key_id: 阿里云访问密钥ID（优先使用环境变量）
        :param access_key_secret: 阿里云访问密钥密钥
        :param workspace_id: 工作空间ID
        :param endpoint: 服务端点
        """
        self.access_key_id = access_key_id or os.getenv('ALIBABA_CLOUD_ACCESS_KEY_ID')
        self.access_key_secret = access_key_secret or os.getenv('ALIBABA_CLOUD_ACCESS_KEY_SECRET')
        self.workspace_id = workspace_id or os.getenv('WORKSPACE_ID')
        self.endpoint = endpoint

class MemoryClient:
    """阿里云百炼记忆服务客户端"""
    
    def __init__(self, config: BailianConfig):
        """
        初始化记忆客户端
        :param config: 配置对象
        """
        if not all([config.access_key_id, config.access_key_secret, config.workspace_id]):
            raise ValueError("Missing required configuration parameters")
            
        self.config = config
        self._init_client()
        
    def _init_client(self):
        """初始化SDK客户端"""
        credential_config = CredentialConfig(
            type="access_key",
            access_key_id=self.config.access_key_id,
            access_key_secret=self.config.access_key_secret
        )
        
        api_config = open_api_models.Config(
            credential=CredentialClient(credential_config),
            endpoint=self.config.endpoint
        )
        
        self.client = BailianClient(api_config)

    def create_memory(self, description: Optional[str] = None) -> Dict[str, Any]:
        """创建新记忆体"""
        """返回示例：
        {
            "memoryId": "6bff4f317a14442fbc9f73d29dbxxxx",
            "requestId": "17204B98-xxxx-4F9A--2446A84821CA"
        }
        """
        request = bailian_models.CreateMemoryRequest(description=description)
        response = self.client.create_memory(
            workspace_id=self.config.workspace_id,
            request=request
        )
        return self._parse_response(response)

    def get_memory(self, memory_id: str) -> Dict[str, Any]:
        """获取记忆体详情"""
        """返回示例：
        {
            "description": "我的大模型应用$APP_ID关于A用户的长期记忆体",
            "memoryId": "6bff4f317a14442fbc9f73d29dbdxxxx\n",
            "requestId": "6a71f2d9-f1c9-913b-818b-11402910xxxx\n",
            "workspaceId": "llm-3z7uw7fwz0vexxxx\n"
        }
        """
        response = self.client.get_memory(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id
        )
        return self._parse_response(response)
    
    def update_memory(self, memory_id: str, description: str) -> Dict[str, Any]:
        """更新记忆体描述"""
        """返回示例：
        {
            "requestId": "6a71f2d9-f1c9-913b-818b-11402910xxxx\n"
        }
        """
        request = bailian_models.UpdateMemoryRequest(description=description)
        response = self.client.update_memory(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            request=request
        )
        return self._parse_response(response)
    
    def delete_memory(self, memory_id: str) -> Dict[str, Any]:
        """删除记忆体"""
        """返回示例：
        {
            "requestId": "6a71f2d9-f1c9-913b-818b-11402910xxxx\n"
        }
        """
        response = self.client.delete_memory(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id
        )
        return self._parse_response(response)
    
    def get_memory_list(self, max_results: int = 10, next_token: Optional[str] = None) -> Dict[str, Any]:
        """获取记忆体列表（单次请求）"""
        """返回示例：
        {
            "maxResults": 10,
            "memories": [
                {
                    "description": "我的大模型应用$APP_ID关于A用户的长期记忆体",
                    "memoryId": "6bff4f317a14442fbc9f73d29dbdxxxx\n"
                }
            ],
            "nextToken": "6a71f2d9-f1c9-913b-818b-11402910xxxx\n",
            "requestId": "6a71f2d9-f1c9-913b-818b-11402910xxxx\n",
            "totalCount": 10,
            "workspaceId": "llm-3z7uw7fwz0vexxxx\n"
        }
        """
        request = bailian_models.ListMemoriesRequest(max_results=max_results, next_token=next_token)
        response = self.client.list_memories(
            workspace_id=self.config.workspace_id,
            request=request
        )
        return self._parse_response(response)
    
    def create_memory_node(self, memory_id: str, content: str) -> Dict[str, Any]:
        """创建记忆片段"""
        """返回示例：
        {
            "memoryNodeId": "68de06c95368463a8be4a84efc872cc5",
            "requestId": "8C56C7AF-6573-19CE-B018-E05E1EDCF4C5"
        }
        """
        request = bailian_models.CreateMemoryNodeRequest(content=content)
        response = self.client.create_memory_node(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            request=request
        )
        return self._parse_response(response)
    
    def get_memory_node(self, memory_id: str, memory_node_id: str) -> Dict[str, Any]:
        """获取记忆片段的内容"""
        """返回示例：
        {
            "content": "用户喜欢吃西红柿炒鸡蛋\n",
            "memoryId": "6bff4f317a14442fbc9f73d29dbd5fc3\n",
            "memoryNodeId": "68de06c95368463a8be4a84efc872cc5",
            "requestId": "8C56C7AF-6573-19CE-B018-E05E1EDCF4C5\n",
            "workspaceId": "llm-us9hjmt32nysdm5v\n"
        }
        """
        response = self.client.get_memory_node(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            memory_node_id=memory_node_id
        )
        return self._parse_response(response)
    
    def update_memory_node(self, memory_id: str, memory_node_id: str, content: str) -> Dict[str, Any]:
        """更新记忆片段的内容"""
        """返回示例：
        {
            "requestId": "8C56C7AF-6573-19CE-B018-E05E1EDCF4C5\n"
        }
        """
        request = bailian_models.UpdateMemoryNodeRequest(content=content)
        response = self.client.update_memory_node(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            memory_node_id=memory_node_id,
            request=request
        )
        return self._parse_response(response)
    
    def delete_memory_node(self, memory_id: str, memory_node_id: str) -> Dict[str, Any]:
        """删除记忆片段"""
        """返回示例：
        {
            "requestId": "8C56C7AF-6573-19CE-B018-E05E1EDCF4C5\n"
        }
        """
        response = self.client.delete_memory_node(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            memory_node_id=memory_node_id
        )
        return self._parse_response(response)
    
    def list_memory_nodes(self, memory_id: str, max_results: int = 10, next_token: Optional[str] = None) -> Dict[str, Any]:
        """获取记忆片段列表"""
        """返回示例：
        {
            "maxResults": 20,
            "memoryNodes": [
                {
                "content": "用户喜欢吃西红柿炒鸡蛋\n",
                "memoryNodeId": "68de06c95368463a8be4a84efc872cc5"
                }
            ],
            "nextToken": "dc270401186b433f975d7e1faaa34e0e\n",
            "requestId": "6a71f2d9-f1c9-913b-818b-114029103cad\n",
            "totalCount": 100
        }
        """
        request = bailian_models.ListMemoryNodesRequest(max_results=max_results, next_token=next_token)
        response = self.client.list_memory_nodes(
            workspace_id=self.config.workspace_id,
            memory_id=memory_id,
            request=request
        )
        return self._parse_response(response)

    def list_all_memory_nodes(self, memory_id: str) -> List[Dict[str, Any]]:
        """获取指定记忆体的全部记忆片段（自动处理分页）"""
        all_nodes = []
        next_token = None

        # 计数器
        count = 0
        
        while True:
            # 每次请求尽可能获取最大数量（API允许的最大值50）
            response = self.list_memory_nodes(
                memory_id=memory_id,
                max_results=50,
                next_token=next_token
            )

            count += 1
            print(f"第 {count} 次请求")
            print(f"response: {response}\n")
            # 添加当前页的记忆节点
            if 'memoryNodes' in response:
                all_nodes.extend(response['memoryNodes'])
            
            # 检查是否有下一页
            next_token = response.get('nextToken')
            if not next_token:
                break
                
        return all_nodes

    def list_all_memories(self) -> List[Dict[str, Any]]:
        """获取全部记忆体（自动分页）"""
        all_memories = []
        next_token = None
        
        while True:
            # 每次获取最大允许数量
            response = self.get_memory_list(
                max_results=50,
                next_token=next_token
            )
            
            # 合并记忆体数据
            if 'memories' in response:
                all_memories.extend(response['memories'])
            
            # 检查分页令牌
            next_token = response.get('nextToken')
            if not next_token:
                break
                
            # 添加API限流保护（1秒/次）
            time.sleep(1)
            
        return all_memories

    def _parse_response(self, response) -> Dict[str, Any]:
        """统一处理响应"""
        if response.status_code != 200:
            raise Exception(f"API请求失败，状态码：{response.status_code}")
        return response.body.to_map()

config = BailianConfig(
    access_key_id=ACCESS_KEY_ID,
    access_key_secret=ACCESS_KEY_SECRET,
    workspace_id=WORKSPACE_ID
)
client = MemoryClient(config)

# 使用示例
if __name__ == "__main__":

    # 已有的记忆体ID
    memory_id = "61501eed73fe4e12984aa38c1382d584"
    
    # client = MemoryClient(config)
    
    # # 创建记忆体
    # new_memory = client.create_memory(description="测试记忆体")
    # print("新建记忆体:", new_memory)
    
    # # 查询记忆体
    # memory_detail = client.get_memory(new_memory['memoryId'])
    # print("记忆体详情:", memory_detail)

    # # 创建记忆片段
    # new_memory_node = client.create_memory_node(memory_id=memory_id, content="用户喜欢吃西红柿炒鸡蛋")
    # print("新建记忆片段:", new_memory_node) # {'memoryNodeId': 'db2b47dbe5d5411a8a7c4a06f9d9a2a3', 'requestId': '9CADADE1-8D1D-52F4-807E-A5D843630064'}

    # # 创建记忆片段
    # new_memory_node = client.create_memory_node(memory_id=memory_id, content="用户喜欢吃西红柿炒鸡蛋")
    # print("新建记忆片段:", new_memory_node) # {'memoryNodeId': 'aa9694421da94646895b32c6e5b4bc1e', 'requestId': '6111C5EC-56C0-58EF-8CD2-B7A4647A0630'}

    # # 获取记忆片段
    # memory_node_detail = client.get_memory_node(memory_id=memory_id, memory_node_id="aa9694421da94646895b32c6e5b4bc1e")
    # print("记忆片段详情:", memory_node_detail)

    # # 获取记忆片段列表
    # memory_node_list = client.list_memory_nodes(memory_id=memory_id, next_token="a0f5d84645cd428f8788a7b5b75676b7")
    # print("记忆片段列表:", memory_node_list)
    
    # # 获取指定记忆体的全部记忆节点
    # all_nodes = client.list_all_memory_nodes(memory_id=memory_id)
    # print(f"共获取到 {len(all_nodes)} 条记忆片段")

    # # 遍历查看内容
    # for node in all_nodes:
    #     print(f"ID: {node['memoryNodeId']}, 内容: {node['content']}")
    
    # 获取全部记忆体
    all_memories = client.list_all_memories()
    print(f"共发现 {len(all_memories)} 个记忆体")

    # 查看记忆体详情
    for memory in all_memories:
        print(f"ID: {memory['memoryId']}, 描述: {memory.get('description', '无描述')}")
    
    